# -*- coding: utf-8 -*-
"""
Created on Wed Jun  1 20:53:07 2022

@author: wjcy19870122
Changed by XPC
"""
import utils
import torch
import torchvision
import transforms as T
from ChmInstanceSegDataset import ChmInstanceSegDataset
from engine import train_one_epoch, evaluate
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor


def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)#使用该模型进行目标检测

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       num_classes)

    return model

def get_model_instance_segmentation_v2(num_classes):
    backbone = torchvision.models.mobilenet_v2(pretrained=True).features
    backbone.out_channels = 1280
    from torchvision.models.detection.rpn import AnchorGenerator
    anchor_generator = AnchorGenerator(sizes=((8, 16, 32, 64, 128),),
                                   aspect_ratios=((0.5, 1.0, 2.0),))
    
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
                                                output_size=7,
                                                sampling_ratio=2)
    
    model = torchvision.models.detection.MaskRCNN(backbone,
                   num_classes=num_classes,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)
    return model


def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    transforms.append(T.Resize((550, 800)))
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5)) 
    #即水平翻转
    return T.Compose(transforms)

def start_train():
    # train on the GPU or on the CPU, if a GPU is not available
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # our dataset has two classes only - background and person 进行二分类
    num_classes = 2
    # use our dataset and defined transformations
    dataset = ChmInstanceSegDataset('data/train', get_transform(train=True))#json格式文件水平翻转
    dataset_test = ChmInstanceSegDataset('data/test', get_transform(train=False))

    '''# split the dataset in train and test set
    indices = torch.randperm(len(dataset)).tolist()
    dataset = torch.utils.data.Subset(dataset, indices[:-50])
    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])'''

    # define training and validation data loaders
    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=4, shuffle=True, num_workers=4,
        collate_fn=utils.collate_fn)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=4, shuffle=False, num_workers=4,
        collate_fn=utils.collate_fn)

    # get the model using our helper function
    model = get_model_instance_segmentation(num_classes)#模型实例分割模型
    #model = get_model_instance_segmentation_v2(num_classes)
    # move model to the right device
    model.to(device)
#
    # construct an optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.005,
                                momentum=0.9, weight_decay=0.0005)
    # and a learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                   step_size=3,
                                                   gamma=0.1)

    # let's train it for 20 epochs
    num_epochs = 20 #先训练五轮,不然容易出问题

    for epoch in range(num_epochs):
        # train for one epoch, printing every 10 iterations
        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
        # update the learning rate
        lr_scheduler.step()
        # evaluate on the test dataset
        evaluate(model, data_loader_test, device=device)

        torch.save(model, 'checkpoints/weights_epoch{0}.pth'.format(epoch))
    print("That's it!")
    
    
if __name__ == '__main__':
    start_train()